import os
import sys
import hashlib
from optparse import OptionParser
from collections import defaultdict

def get_files_size(directory, min_size):
    """读取指定目录下所有文件的大小，返回大小到文件路径列表的字典"""
    if not os.path.isdir(directory):
        print(f"错误: 目录 '{directory}' 不存在或不是一个有效的目录")
        return None
    
    file_sizes = defaultdict(list)
    
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            try:
                file_size = os.path.getsize(file_path)
                if file_size >= min_size:
                    file_sizes[file_size].append(file_path)
            except OSError as e:
                print(f"无法获取文件 '{file_path}' 的大小: {e}")
    
    return file_sizes

def get_filename(file_path):
    """提取文件名（不含路径）"""
    return os.path.basename(file_path)

def get_md5_hash(file_path, block_size=4096):
    """计算文件前block_size字节的MD5哈希值"""
    md5 = hashlib.md5()
    try:
        with open(file_path, 'rb') as f:
            data = f.read(block_size)  # 只读取前block_size字节
            if data:
                md5.update(data)
        return md5.hexdigest()
    except OSError as e:
        print(f"无法读取文件 '{file_path}' 计算MD5: {e}")
        return None

def calculate_statistics(file_sizes):
    """计算并返回文件统计信息"""
    stats = {
        'total_files': 0,
        'single_file_groups': 0,
        'multiple_file_groups': 0
    }
    
    for size in file_sizes:
        files_count = len(file_sizes[size])
        stats['total_files'] += files_count
        if files_count == 1:
            stats['single_file_groups'] += 1
        else:
            stats['multiple_file_groups'] += 1
    
    return stats

def find_duplicate_groups(files, read_bytes):
    """查找并返回重复文件组（同名或同MD5）"""
    duplicate_groups = defaultdict(list)
    name_to_md5 = {}
    
    # 预先计算所有文件的MD5
    for file_path in files:
        name_to_md5[file_path] = get_md5_hash(file_path, read_bytes)
    
    # 1. 处理同名文件
    name_groups = defaultdict(list)
    for file_path in files:
        name = get_filename(file_path)
        name_groups[name].append(file_path)
    
    processed_files = set()
    for name, paths in name_groups.items():
        if len(paths) > 1:
            # 使用第一个文件的MD5作为组标识，或生成一个基于名称的标识
            group_id = name_to_md5[paths[0]] or f"name:{name}"
            duplicate_groups[group_id].extend(paths)
            processed_files.update(paths)
    
    # 2. 处理剩余文件（不同名），按MD5分组
    remaining_files = [f for f in files if f not in processed_files]
    md5_groups = defaultdict(list)
    for file_path in remaining_files:
        md5_hash = name_to_md5[file_path]
        if md5_hash:
            md5_groups[md5_hash].append(file_path)
    
    for md5_hash, paths in md5_groups.items():
        if len(paths) > 1:
            duplicate_groups[md5_hash].extend(paths)
    
    return duplicate_groups

def get_user_selection(paths):
    """获取用户选择的要保留的文件"""
    # 提示用户操作方式
    print("    请选择要保留的文件:")
    print("    - 直接回车: 默认保留第1个文件")
    print("    - 输入序号: 保留对应文件")
    print("    - 输入n: 取消操作")
    print("    - 输入a: 删除全部")
    
    while True:
        answer = input(f"    请输入选择 (1-{len(paths)}): ").strip().lower()
        
        # 处理直接回车（空输入）的情况，默认保留第1个
        if not answer:
            return paths[0], paths[1:]
        
        # 处理取消操作
        if answer == 'n':
            print("    已取消删除操作")
            return None, []
        
        # 处理取消操作
        if answer == 'a':
            print("    删除此组全部文件")
            return None, paths
        
        # 处理数字输入
        try:
            selected = int(answer)
            if 1 <= selected <= len(paths):
                keep_index = selected - 1  # 转换为列表索引
                keep = paths[keep_index]
                to_delete = [p for idx, p in enumerate(paths) if idx != keep_index]
                return keep, to_delete
            else:
                print(f"    输入无效，请输入1到{len(paths)}之间的数字")
        except ValueError:
            print("    输入无效，请输入数字、n或直接回车")

def process_duplicate_groups(duplicate_groups, auto_delete):
    """处理重复文件组，根据auto_delete参数决定是否自动删除"""
    deleted_count = 0
    
    if duplicate_groups:
        print("\n该组中发现以下重复文件组：")
        
        for i, (group_id, paths) in enumerate(duplicate_groups.items(), 1):
            # 确定组类型描述
            if group_id.startswith("name:"):
                group_type = f"同名文件组（文件名: {group_id[5:]}）"
            else:
                group_type = f"内容相同文件组（MD值一样: {group_id}）"
            
            print(f"  重复组 {i}: {group_type}（共 {len(paths)} 个文件）")
            for j, path in enumerate(paths, 1):
                print(f"    {j}. {path}")
            
            # 处理删除逻辑
            if auto_delete:
                # 自动删除模式，保留第一个文件
                keep = paths[0]
                to_delete = paths[1:]
            else:
                # 交互模式，让用户选择要保留的文件
                keep, to_delete = get_user_selection(paths)
            
            # 执行删除
            if to_delete:
                print(f"    保留: {keep}")
                for path in to_delete:
                    try:
                        os.remove(path)
                        print(f"    已删除: {path}")
                        deleted_count += 1
                    except OSError as e:
                        print(f"    无法删除 {path}: {e}")
        print()  # 空行分隔
    
    return deleted_count

def print_statistics(stats, duplicate_groups_count, deleted_files):
    """打印最终统计信息"""
    print("统计结果：")
    print(f"1. 总文件数量：{stats['total_files']}")
    print(f"2. 单文件大小组数量：{stats['single_file_groups']}")
    print(f"3. 多文件大小组数量：{stats['multiple_file_groups']}")
    print(f"4. 重复文件组总数（同名或内容相同）：{duplicate_groups_count}")
    print(f"5. 已删除的重复文件数量：{deleted_files}")

def main():
    # 解析命令行参数
    parser = OptionParser(usage="usage: %prog [options] directory")
    parser.add_option("-y", "--yes", "--auto-delete",
                      action="store_true", dest="auto_delete", default=False,
                      help="自动删除重复文件，无需交互确认 [默认: 关闭]")
    parser.add_option("-b", "--bytes", type="int", dest="read_bytes", 
                      default=-1, metavar="NUM",
                      help="计算MD5时读取的字节数，-1表示读取整个文件 [默认: 4096]")
    parser.add_option("-m", "--min_size", type="int", dest="min_read_bytes", 
                      default=0, metavar="NUM",
                      help="统计文件最小值 [默认: 0]")
    
    (options, args) = parser.parse_args()
    
    # 验证目录参数
    if not args:
        parser.error("请指定要扫描的目录")
    
    # 验证读取字节数的有效性
    if options.read_bytes < -1:
        parser.error(f"无效的字节数: {options.read_bytes}，必须是-1或正整数")
    
    # 确定目标目录
    target_directory = args[0]
    print(f"正在读取目录 '{target_directory}' 下的所有文件大小, md5_size:{options.read_bytes}...\n")
    
    # 获取文件大小分组
    file_sizes = get_files_size(target_directory, options.min_read_bytes)
    if not file_sizes:
        return
    
    # 计算基本统计信息
    stats = calculate_statistics(file_sizes)
    
    # 处理多文件大小组
    deleted_files = 0
    duplicate_groups_count = 0
    
    if stats['multiple_file_groups'] > 0:
        print("发现以下具有相同大小的文件组：\n")
        
        for size in sorted(file_sizes.keys()):
            files = file_sizes[size]
            if len(files) > 1:
                # 打印该大小组的所有文件
                print(f"文件大小: {size} 字节（共 {len(files)} 个文件）")
                for i, path in enumerate(files, 1):
                    print(f"  {i}. {path}")
                
                # 查找重复文件组
                duplicate_groups = find_duplicate_groups(files, options.read_bytes)
                duplicate_groups_count += len(duplicate_groups)
                
                # 处理重复文件组并统计删除数量
                deleted_files += process_duplicate_groups(duplicate_groups, options.auto_delete)
                
                if not duplicate_groups:
                    print("该组中没有重复文件，不执行删除操作\n")
    
    # 打印最终统计
    print_statistics(stats, duplicate_groups_count, deleted_files)

if __name__ == "__main__":
    main()
    